from typing import Dict, Any, TypedDict, Annotated
from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage, AIMessage
import asyncio
import time

class ParallelState(TypedDict):
    messages: Annotated[list, "对话消息列表"]
    result: Annotated[str, "处理结果"]
    parallel_results: Annotated[Dict[str, Any], "并行处理结果"]
    combined_result: Annotated[str, "合并结果"]

class ParallelGraphDemo:
    """并行执行图演示"""
    
    def __init__(self):
        self.graph = self._build_graph()
    
    def _build_graph(self):
        """构建并行执行图"""
        
        workflow = StateGraph(ParallelState)
        
        # 添加节点
        workflow.add_node("prepare_tasks", self._prepare_tasks)
        workflow.add_node("task_a", self._task_a)
        workflow.add_node("task_b", self._task_b)
        workflow.add_node("task_c", self._task_c)
        workflow.add_node("combine_results", self._combine_results)
        workflow.add_node("format_output", self._format_output)
        
        # 设置入口点
        workflow.set_entry_point("prepare_tasks")
        
        # 设置并行边
        workflow.add_edge("prepare_tasks", "task_a")
        workflow.add_edge("prepare_tasks", "task_b")
        workflow.add_edge("prepare_tasks", "task_c")
        
        # 设置合并边
        workflow.add_edge("task_a", "combine_results")
        workflow.add_edge("task_b", "combine_results")
        workflow.add_edge("task_c", "combine_results")
        
        workflow.add_edge("combine_results", "format_output")
        workflow.add_edge("format_output", END)
        
        return workflow.compile()
    
    def _prepare_tasks(self, state: ParallelState) -> ParallelState:
        """准备并行任务"""
        messages = state["messages"]
        if not messages:
            messages = [HumanMessage(content="执行并行任务")]
        
        return {
            "messages": messages,
            "result": "任务已准备",
            "parallel_results": {},
            "combined_result": ""
        }
    
    def _task_a(self, state: ParallelState) -> ParallelState:
        """任务A - 模拟数据处理"""
        time.sleep(0.1)  # 模拟处理时间
        
        return {
            "messages": state["messages"],
            "result": state["result"],
            "parallel_results": {"task_a": "数据处理完成"},
            "combined_result": ""
        }
    
    def _task_b(self, state: ParallelState) -> ParallelState:
        """任务B - 模拟网络请求"""
        time.sleep(0.2)  # 模拟网络延迟
        
        return {
            "messages": state["messages"],
            "result": state["result"],
            "parallel_results": {"task_b": "网络请求完成"},
            "combined_result": ""
        }
    
    def _task_c(self, state: ParallelState) -> ParallelState:
        """任务C - 模拟文件操作"""
        time.sleep(0.15)  # 模拟文件操作时间
        
        return {
            "messages": state["messages"],
            "result": state["result"],
            "parallel_results": {"task_c": "文件操作完成"},
            "combined_result": ""
        }
    
    def _combine_results(self, state: ParallelState) -> ParallelState:
        """合并并行结果"""
        # 收集所有并行任务的结果
        all_results = {}
        
        # 这里需要从多个输入状态中收集结果
        # 在实际的LangGraph中，这通常通过join操作完成
        if "task_a" in state.get("parallel_results", {}):
            all_results["task_a"] = state["parallel_results"]["task_a"]
        if "task_b" in state.get("parallel_results", {}):
            all_results["task_b"] = state["parallel_results"]["task_b"]
        if "task_c" in state.get("parallel_results", {}):
            all_results["task_c"] = state["parallel_results"]["task_c"]
        
        combined = f"合并结果: {all_results}"
        
        return {
            "messages": state["messages"],
            "result": state["result"],
            "parallel_results": all_results,
            "combined_result": combined
        }
    
    def _format_output(self, state: ParallelState) -> ParallelState:
        """格式化输出"""
        return {
            "messages": state["messages"],
            "result": f"并行执行完成: {state['combined_result']}",
            "parallel_results": state["parallel_results"],
            "combined_result": state["combined_result"]
        }
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行并行执行图"""
        initial_state = {
            "messages": [HumanMessage(content=input_data.get("message", "执行并行任务"))],
            "result": "",
            "parallel_results": {},
            "combined_result": ""
        }
        
        result = self.graph.invoke(initial_state)
        
        return {
            "final_state": result,
            "parallel_results": result["parallel_results"],
            "combined_result": result["combined_result"],
            "final_result": result["result"]
        }
